#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static const char base64_table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

/**
 * @brief 编码输入的字节数组为Base64字符串
 *
 * @param input 输入的字节数组
 * @param length 输入字节数组的长度
 * @return 编码后的Base64字符串（调用者需要释放内存），
 *         如果内存分配失败，则返回NULL。
 */
char* base64_encode(const unsigned char *input, int length) {
    int output_length = 4 * ((length + 2) / 3); // 计算输出长度
    char *output = (char *)malloc(output_length + 1); // 动态分配内存

    if (output == NULL) {
        return NULL; // 内存分配失败
    }

    for (int i = 0, j = 0; i < length;) {
        unsigned int octet_a = (i < length) ? input[i++] : 0;
        unsigned int octet_b = (i < length) ? input[i++] : 0;
        unsigned int octet_c = (i < length) ? input[i++] : 0;

        unsigned int triple = (octet_a << 16) + (octet_b << 8) + octet_c;

        output[j++] = base64_table[(triple >> 18) & 0x3F];
        output[j++] = base64_table[(triple >> 12) & 0x3F];
        output[j++] = (i > length + 1) ? '=' : base64_table[(triple >> 6) & 0x3F];
        output[j++] = (i > length) ? '=' : base64_table[triple & 0x3F];
    }
    output[output_length] = '\0'; // 以空字符结束字符串
    return output; // 返回编码后的字符串
}

/**
 * @brief 解码Base64字符串为原始字节数组
 *
 * @param input 输入的Base64字符串
 * @param output 输出的字节数组
 * @return 解码后的字节数组长度，
 *         如果输入字符串无效（如包含非法字符），则返回-1。
 */
int base64_decode(const char *input, unsigned char *output) {
    int input_length = strlen(input);
    unsigned char temp[4];
    int length = 0;

    for (int i = 0; i < input_length;) {
        for (int j = 0; j < 4 && i < input_length; j++, i++) {
            if (input[i] == '=') {
                temp[j] = 0;
            } else {
                char *ptr = strchr(base64_table, input[i]);
                if (ptr == NULL) {
                    return -1; // 输入字符串无效
                }
                temp[j] = ptr - base64_table;
            }
        }

        output[length++] = (temp[0] << 2) + (temp[1] >> 4);
        if (temp[2] != 0) {
            output[length++] = (temp[1] << 4) + (temp[2] >> 2);
        }
        if (temp[3] != 0) {
            output[length++] = (temp[2] << 6) + temp[3];
        }
    }
    return length; // 返回解码后的长度
}

int main() {
    const char *text = "Hello, World!";
    char *encoded = base64_encode((unsigned char *)text, strlen(text));
    if (encoded == NULL) {
        fprintf(stderr, "Memory allocation failed\n");
        return 1;
    }

    printf("Encoded: %s\n", encoded);

    unsigned char decoded[256];
    int decoded_length = base64_decode(encoded, decoded);
    if (decoded_length < 0) {
        fprintf(stderr, "Invalid Base64 input\n");
        free(encoded);
        return 1;
    }
    decoded[decoded_length] = '\0'; // Null-terminate the decoded string
    printf("Decoded: %s\n", decoded);

    free(encoded); // 释放动态分配的内存
    return 0;
}
